package com.liangchao.androidmediacodectest;

import android.content.Context;
import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaCodecList;
import android.media.MediaFormat;
import android.util.Log;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;
import java.util.Calendar;
import java.util.Vector;

public class EncodeUtil implements Runnable{

    private final static String TAG = "EncodeUtil";

    private String mimeType;
    private MediaCodec mediaCodec;
    private MediaCodecInfo selectedCodecInfo;
    private MediaFormat mediaFormat;
    private int mInputSource = 0;
    //raw数据中提取yuv图像的张数
    private final int YUV_NUM_IN_RAW=9;
    //用于写编码后文件
    private FileOutputStream fileOutputStream = null;

    private InputStream inputStream;

    private Context mContext;
    private boolean mStopFlag = false;
    private int outCount;
    private int timeTotal = 0;
    private EncodeProgressListener encodeProgressListener;


    private int INPUT_WIDTH,INPUT_HEIGHT;

    public void setINPUT_WIDTH(int INPUT_WIDTH) {
        this.INPUT_WIDTH = INPUT_WIDTH;
    }

    public void setINPUT_HEIGHT(int INPUT_HEIGHT) {
        this.INPUT_HEIGHT = INPUT_HEIGHT;
    }
    public EncodeUtil(Context context){
        mContext = context;
    }

    public void  setMimeType(String mimeType){
        this.mimeType = mimeType;
    }

    public void setInputSource(int source){
        mInputSource= source;
    }

    public void setEncodeProgressListener(EncodeProgressListener encodeProgressListener){
        this.encodeProgressListener = encodeProgressListener;
    }

    public void initEncoder (){
        if(mimeType.isEmpty()){
            Log.e(TAG,"please set the mimeType");
        }
        selectedCodecInfo = selectCodecInfo();

        mediaFormat = MediaFormat.createVideoFormat(mimeType, INPUT_WIDTH, INPUT_HEIGHT);//这里我们需要指定mimeType以及编码后的宽高
        mediaFormat = MediaFormat.createVideoFormat(mimeType, INPUT_WIDTH, INPUT_HEIGHT);
//mediaFormat.setInteger(MediaFormat.KEY_COLOR_FORMAT, MediaCodecInfo.CodecCapabilities.COLOR_QCOM_FormatYUV420SemiPlanar);
        mediaFormat.setInteger(MediaFormat.KEY_COLOR_FORMAT, 21);
        mediaFormat.setInteger(MediaFormat.KEY_BIT_RATE, 4 * 1024 * 1024*1024);
        mediaFormat.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, 4);
        mediaFormat.setInteger(MediaFormat.KEY_BITRATE_MODE, MediaCodecInfo.EncoderCapabilities.BITRATE_MODE_CBR);
        mediaFormat.setInteger(MediaFormat.KEY_FRAME_RATE, 30); //帧率
        //mediaFormat.setInteger("vendor.qti-ext-enc-colorspaceconversion.enable",1);//高通颜色转换bt709隐藏参数
        mediaFormat.setInteger ("vendor.qti-ext-enc-bitrate-mode.value",2);//这里为高通特有参数，我们可以根据不同的芯片平台设置他们独有的功能（需要参考平台文档），比如这里就可以指定码率的模式
        try {
            mediaCodec = MediaCodec.createByCodecName(selectedCodecInfo.getName());
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        mediaCodec.configure(mediaFormat,null,null,MediaCodec.CONFIGURE_FLAG_ENCODE);
        mediaCodec.start();
        String savePath = mContext.getExternalFilesDir(null)+"/"+"encode.h264";
        try {
            fileOutputStream = new FileOutputStream(new File(savePath));
        } catch (FileNotFoundException e) {
            throw new RuntimeException(e);
        }

    }

    private MediaCodecInfo selectCodecInfo(){
        int codecNum = MediaCodecList.getCodecCount();
        Log.d(TAG,"codecNum:"+codecNum);

        for(int i=0; i<codecNum;i++){
            MediaCodecInfo mediaCodecInfo = MediaCodecList.getCodecInfoAt(i);
            Log.d(TAG,"codec info "+mediaCodecInfo.getName());
            if(!mediaCodecInfo.isEncoder()){
                continue;
            }
            String[] types = mediaCodecInfo.getSupportedTypes();
            for(int j = 0; j< types.length; j++){
                if(types[j].equalsIgnoreCase(mimeType)){
                    return  mediaCodecInfo;
                }
            }
        }
        Log.e(TAG,"Cant get right MediaInfo");
        return null;
    }

    @Override
    public void run() {
        Log.d(TAG,"Now will start encode");

        //解码后获取buffer元数据信息
        MediaCodec.BufferInfo bufferInfo = new MediaCodec.BufferInfo();

        long timeoutUs = 12*1000;
        long mEndTime = 0L;
        long curretTime = 0L;
        boolean isFirstFrame = true;

        ByteBuffer[] inputBuffers = mediaCodec.getInputBuffers();
        ByteBuffer[] outputBuffers = mediaCodec.getOutputBuffers();

        if(mInputSource != 0){
            inputStream = mContext.getResources().openRawResource(mInputSource);

        }


        //拆分inputStream 将每张yuv数据储存到向量数组里
        Vector<byte[]> bytesInputYUV = new Vector<>();
        //提取yuv图像
        try {
            for(int i=0; i < YUV_NUM_IN_RAW; i++){
                byte[] oneYUVPic = new byte[INPUT_WIDTH*INPUT_HEIGHT*3/2];
                inputStream.read(oneYUVPic,0,INPUT_WIDTH*INPUT_HEIGHT*3/2);
                bytesInputYUV.add(oneYUVPic);
            }

            inputStream.close();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }

        Calendar calendar = Calendar.getInstance();
        long startMils = calendar.getTimeInMillis();
        Log.d(TAG,"Start time "+calendar.getTime());
        int inputCount =0;
        long startTime = System.nanoTime();
        long encodeTime;
        while (!mStopFlag){

            int inputBufferId = mediaCodec.dequeueInputBuffer(timeoutUs);
            //准备数据并且将其dequeue到inputBuffer
            if(inputBufferId >= 0){
                ByteBuffer inputBuffer = inputBuffers[inputBufferId];
                if(inputCount++ == 2700)//编码满2700后结束
                    break;
                float progressEncode = (float)inputCount/(float) 2700;
                int porgress =(int) (progressEncode*100);
                encodeProgressListener.onEncodeProgress(porgress);
                byte[] data = bytesInputYUV.get(inputCount%YUV_NUM_IN_RAW);
                inputBuffer.put(data,0,data.length);
                encodeTime = (System.nanoTime() - startTime)/1000+100;
                mediaCodec.queueInputBuffer(inputBufferId,0,data.length,encodeTime,0);
                inputBuffer.clear();
            }


            int outpueBufferId = mediaCodec.dequeueOutputBuffer(bufferInfo,timeoutUs);
            if(outpueBufferId >= 0){
                startTime = System.currentTimeMillis();
                ByteBuffer outputBuffer = outputBuffers[outpueBufferId];
                byte[] outStream = new byte[bufferInfo.size];
                outputBuffer.get(outStream,0,bufferInfo.size);
                try {
                    fileOutputStream.write(outStream);
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }

                outCount++;
                if(outCount%60 == 0){
                    Log.d(TAG,"inputCount : "+inputCount+" outCount : "+outCount);
                }
                mediaCodec.releaseOutputBuffer(outpueBufferId,0);
                curretTime = System.currentTimeMillis();
                if(isFirstFrame){
                    isFirstFrame = false;
                    mEndTime = curretTime;
                }else {
                    long elapsedTime = curretTime - mEndTime;
                    timeTotal++;
                }
            }

        }
        try {
            mediaCodec.stop();
            mediaCodec.release();
            fileOutputStream.close();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }

    }
    public interface EncodeProgressListener{
        void onEncodeProgress(int progress);
    }

}
